from aiogram import Router, types
from aiogram.filters.command import Command
from pymongo import MongoClient

from datetime import datetime, timedelta
import json
from decouple import config

# Your MongoDB connection setup
client = MongoClient(config("CLIENT"))
db = client[config("DB")]
collection = db[config("COLLECTION")]

router = Router()

# Handler for /start command
@router.message(Command(commands=['start']))
async def send_welcome(message: types.Message):
    await message.reply("Привет! Это бот для агрегации данных.")
    await message.reply("Отправь мне JSON-сообщение в формате:\n\n"
                        "```\n"
                        "{\n"
                        "    \"dt_from\": \"2022-02-01T00:00:00\",\n"
                        "    \"dt_upto\": \"2022-02-02T00:00:00\",\n"
                        "    \"group_type\": \"hour\"\n"
                        "}\n"
                        "```")

# Handler for /get_data command
@router.message(Command(commands=['get_data']))
async def aggregate_data_command(message: types.Message):
    await message.reply("Отправь мне JSON-сообщение в формате:\n\n"
                        "```\n"
                        "{\n"
                        "    \"dt_from\": \"2022-02-01T00:00:00\",\n"
                        "    \"dt_upto\": \"2022-02-02T00:00:00\",\n"
                        "    \"group_type\": \"hour\"\n"
                        "}\n"
                        "```")

# Handler for receiving input data
@router.message()
async def aggregate_data(message: types.Message):
    try:
        # Parse JSON message
        data = json.loads(message.text)
        
        # Extract input data
        dt_from = datetime.fromisoformat(data["dt_from"])
        dt_upto = datetime.fromisoformat(data["dt_upto"])
        group_type = data["group_type"]
        
        # Query data from MongoDB
        data = collection.find({
            'dt': {
                '$gte': dt_from,
                '$lt': dt_upto
            }
        }, {'_id': 0, 'value': 1, 'dt': 1})
        
        # Perform data aggregation
        # Your aggregation logic here

        data_list = list(data)
    
    # Если данных нет, вернуть пустые списки
        if not data_list:
            return JSONResponse(content={"dataset": [], "labels": []})
        
        # Функция для получения начала следующего периода агрегации
        def get_next_period_start(date, group_type):
            if group_type == 'hour':
                return date + timedelta(hours=1)
            elif group_type == 'day':
                return date + timedelta(days=1)
            elif group_type == 'month':
                if date.month == 12:
                    return date.replace(year=date.year + 1, month=1, day=1, hour=0, minute=0, second=0)
                else:
                    return date.replace(month=date.month + 1, day=1, hour=0, minute=0, second=0)
            else:
                raise ValueError("Invalid group_type. Use 'hour', 'day', or 'month'")
        
        # Агрегация данных
        aggregated_data = []
        labels = []
        current_period_start = dt_from
        current_period_end = get_next_period_start(current_period_start, group_type)
        current_sum = 0
        
        for entry in data_list:
            while entry['dt'] >= current_period_end:
                aggregated_data.append(current_sum)
                labels.append(current_period_start.isoformat())
                current_period_start = current_period_end
                current_period_end = get_next_period_start(current_period_start, group_type)
                current_sum = 0
            current_sum += entry['value']
        
        # Добавление последнего периода
        aggregated_data.append(current_sum)
        labels.append(current_period_start.isoformat())
        
        # Reply with aggregated data
        await message.reply("Агрегированные данные:\n\n"
                        "```\n"
                        "{\n"
                        f"    \"dataset\": \"{aggregated_data}\",\n"
                        f"    \"labels\": \"{labels}\"\n"
                        "}\n"
                        "```")
        
    except Exception as e:
        await message.reply(f"Error: {str(e)}")